import json
from datetime import datetime


def v1(messages: list, tools: list) -> str:
    """
    :param messages: role+content格式，role包括
    :param tools: openapi标准的函数描述

    生成一个符合qwen查询模板的prompt
    """

    if messages[0]["role"]=="system":
        system_prompt = messages[0]["content"]
    else:
        system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."


    curren_date = datetime.now().strftime("%Y-%m-%d")
    tools_prompt = "\n".join([json.dumps(tool,ensure_ascii=False) for tool in tools])
    prompt=f"""<|im_start|>system

Current Date: {curren_date}
{system_prompt}

# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{tools_prompt}
<tools>

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{{"name": <function-name>, "arguments": <args-json-object>}}
</tool_call><|im_end|>
"""
    for msg in messages:
        if "tool_calls" in msg:
            tool_prompt="\n".join(["<tool_call>\n"+tool_call+"\n</tool_call>" for tool_call in msg['tool_calls']])
            text = "<|im_start|>assistant\n" +tool_prompt+ "<|im_end|>\n"
        elif msg['role']=='user' or msg['role']=='assistant':
            text = "<|im_start|>" + msg['role'] + "\n" + msg['content'] + "<|im_end|>\n"
        else:
            raise Exception("Invalid message role: %s" % msg['role'])

        prompt += text
    return prompt


def v2(messages: list, tools: list=None) -> str:
    """
    :param messages: role+content格式，role包括
    :param tools: openapi标准的函数描述

    生成一个符合qwen查询模板的prompt
    """

    if messages[0]["role"]=="system":
        system_prompt = messages[0]["content"]
    else:
        system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
    curren_date = datetime.now().strftime("%Y-%m-%d")

    if tools:
        tools_prompt = "\n".join([json.dumps(tool,ensure_ascii=False) for tool in tools])
        prompt=f"""<|im_start|>System

Current Date: {curren_date}
{system_prompt}

# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{tools_prompt}
<tools>

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{{"name": <function-name>, "arguments": <args-json-object>}}
</tool_call><|im_end|>
"""
    else:
        prompt = f"""<|im_start|>system

Current Date: {curren_date}
{system_prompt}<|im_end|>
"""


    for msg in messages:
        if "tool_calls" in msg:
            tool_prompt="\n".join(["<tool_call>\n"+tool_call+"\n</tool_call>" for tool_call in msg['tool_calls']])
            text = "<|im_start|>Assistant:\n" +tool_prompt+ "<|im_end|>\n"
        elif msg['role']=='assistant':
            text = "<|im_start|>Assistant:\n" + msg['content'] + "<|im_end|>\n"
        elif msg['role'] == 'user':
            text = "<|im_start|>User:\n" + msg['content'] + "<|im_end|>\n"
        elif msg['role']=='system':
            continue
        else:
            raise Exception("Invalid message role: %s" % msg['role'])

        prompt += text
    prompt+="<|im_start|>Assistant:"
    return prompt


if __name__ == '__main__':
    messages=[
        {"role": "user", "content": "今天天气怎么样？"},
        {"role": "assistant", "tool_calls": [
            {"function": {"name": "get_current_temperature", "arguments": {"location": "San Francisco, CA, USA"}}}
        ]},
        {"role": "user", "content": "北京天气：晴天，15°c"},
        {"role": "assistant", "content": "今天北京是晴天，气温15°c。"}
    ]
    tools=[
        {"type": "function",
         "function": {"name": "get_current_temperature", "description": "Get current temperature at a location.",
                      "parameters": {"type": "object", "properties": {"location": {"type": "string",
                                                                                   "description": "The location to get the temperature for, in the format \"City, State, Country\"."},
                                                                      "unit": {"type": "string",
                                                                               "enum": ["celsius", "fahrenheit"],
                                                                               "description": "The unit to return the temperature in. Defaults to \"celsius\"."}},
                                     "required": ["location"]}}}
    ]

    prompt=v1(messages,tools)
    print(prompt)

